from collections import defaultdict
import json
import re

from model import deepseek, qwen_72B
import concurrent.futures
import threading
from latex2sympy2 import latex2sympy


def read(file_path):
    data = defaultdict(dict)
    with open(file_path, 'r', encoding='utf-8') as file:
        lines = file.readlines()
    
    for index, line in enumerate(lines):
        tmp = json.loads(line)
        data[tmp['key']] = tmp
    return data

def convert_to_float(s):
    s = s.strip()
    try:
        return float(s)
    except:
        try:
            return float(latex2sympy(s))
        except:
            try:
                s = re.sub(r'\\boxed{(.*)}', r'\1', s, flags=re.DOTALL)
                return float(latex2sympy(s))
            except:
                try:
                    s = re.sub(r'\boxed{(.*)}', r'\1', s, flags=re.DOTALL)
                    return float(latex2sympy(s))
                except:
                    return None

def process(data, ori_tag, tgt_tag):
    lock = threading.Lock()
    
    def process_function(key, txt):
        try:
            is_float = convert_to_float(txt)
            if not is_float:
                prompt = f'原文："{txt}"\n任务：将原文转换为标准阿拉伯数字形式。如果不能转换，请返回空字符串。\n\n示例：\n输入：sixty three point零七四\n输出：63.074\n输入：plus three equals\n输出：""\n\n要求：\n- 仅输出转换后的数字，不添加任何说明或解释；\n- 结果必须与原文所表达的数值完全一致。'
                while True:
                    try:
                        n_txt = deepseek(prompt)
                        break
                    except:
                        pass
                print(rf'[{txt}] -> [{n_txt}]')
            else:
                n_txt = txt
            with lock:
                data[key][tgt_tag] = n_txt
        except Exception as e:
            print(f'{key} error')
    
    with concurrent.futures.ThreadPoolExecutor(max_workers=64) as executor:
        futures = []
        for key, value in data.items():
            txt = value[ori_tag]
            futures.append(executor.submit(process_function, key, txt))
        
        for future in concurrent.futures.as_completed(futures):
            pass  # 异常已在process函数中处理

    return data

def write(file_path, error_path, data, tgt_tag):
    with open(file_path, 'w', encoding='utf-8') as file, open(error_path, 'w', encoding='utf-8') as error:
        for key, value in data.items():
            if tgt_tag in value.keys():
                json.dump(value, file, ensure_ascii=False)
                file.write('\n')
            else:
                json.dump(value, error, ensure_ascii=False)
                error.write('\n')

if __name__ == '__main__':
    file_path = '/mnt/sfs/asr/code/osum_xlgeng_3B/voicebench_wsy/qwen2_omni_chat/1.jsonl.n'
    out_path = file_path + '.n'
    err_path = file_path + '.e'
    ori_tag = 'predict_answer'
    tgt_tag = 'predict_answer_itn'

    data = read(file_path)
    data_n = process(data, ori_tag, tgt_tag)
    write(out_path, err_path, data_n, tgt_tag)